import numpy as np
import cv2
from .cal_area import cal_area
# 计算两个轮廓的overlap和union
def cal_overlap_union(cnt_1,cnt_2):
    image_1 = np.zeros((int(np.array(cnt_1+cnt_2).flatten().max())+1,int(np.array(cnt_1+cnt_2).flatten().max())+1),np.uint8)
    image_2 = image_1.copy()
    cv2.drawContours(image_1, contours = [np.int0(cnt_1)], contourIdx = 0, color = 1, thickness=-1)
    cv2.drawContours(image_2, contours = [np.int0(cnt_2)], contourIdx = 0, color = 1, thickness=-1)
    image = image_1+image_2
    # 稍微处理一下，方便数像素
    image = image.flatten().tolist()
    overlap = image.count(2)
    union = image.count(1)+image.count(2)

    return overlap,union
# 计算两组不同层的片构成的3D ct数据的交并比
def cal_IoU_set1_set2(set1,set2):
    layer_set = list(set(list(set1.keys())+list(set2.keys())))
    overlaps = []
    unions = []
    for i in layer_set:
        if i in set1.keys() and i in set2.keys():
            overlap,union = cal_overlap_union(set1[i],set2[i])
            overlaps.append(overlap)
            unions.append(union)
        else:
            overlaps.append(0)
            if(i in set1.keys()):
                unions.append(cal_area(set1[i]))
            if(i in set2.keys()):
                unions.append(cal_area(set2[i]))
    return sum(overlaps)/sum(unions)
# 计算两组不同层的片构成的3D ct数据的Dice
def cal_dice_set1_set2(set1,set2):
    layer_set = list(set(list(set1.keys())+list(set2.keys())))
    overlaps = []
    unions = []
    for i in layer_set:
        if i in set1.keys() and i in set2.keys():
            overlap,union = cal_overlap_union(set1[i],set2[i])
            overlaps.append(overlap)
            unions.append(union)
        else:
            overlaps.append(0)
            if(i in set1.keys()):
                unions.append(cal_area(set1[i]))
            if(i in set2.keys()):
                unions.append(cal_area(set2[i]))
    # ss = []
    # for i in range(len(overlaps)):
    #     ss.append(2*overlaps[i]/(unions[i]+overlaps[i]))
    # return np.mean(ss)
    return 2*sum(overlaps)/(sum(unions)+sum(overlaps))
def cal_IoU(box_dict_of_TP):
    IoU_dict = {}
    for TP_info in box_dict_of_TP.keys():
        IoU_dict[TP_info] = cal_IoU_set1_set2(set1=box_dict_of_TP[TP_info]['gt_boxes'],set2=box_dict_of_TP[TP_info]['ai_boxes'])
    return IoU_dict
def cal_dice(box_dict_of_TP):
    dice_dict = {}
    for TP_info in box_dict_of_TP.keys():
        dice_dict[TP_info] = cal_dice_set1_set2(set1=box_dict_of_TP[TP_info]['gt_boxes'],set2=box_dict_of_TP[TP_info]['ai_boxes'])
    return dice_dict